package Demo;

//
// Source code recreated from a .class file by IntelliJ IDEA
// (powered by Fernflower decompiler)
//

import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.*;

import java.math.BigInteger;
import java.util.*;

public class OP {
    Random random;
    String embedError = "";
    //String extrError = "";
    JarFileLoader loadJar = new JarFileLoader();

    public OP() {
    }

    public Boolean embed(String path, String outpath, String wm, int key) {
        new Fenjie();
        JavaClass jClass = null;
        ArrayList<JavaClass> jClassList = this.loadJar.inputJar(path);
        this.random = new Random();
        this.random.setSeed((long)key);
        ArrayList<Integer> watermark = this.splitWM(wm, String.valueOf(key));
        ArrayList candidates = this.preprocess(jClassList);
        if (watermark.size() > candidates.size()) {
            this.embedError = "the watermark is too long for the application";
            return false;
        } else {
            int i = 0;

            while(i < watermark.size()) {
                Bundle b = this.getBundle(candidates, this.random);
                int ifIndex = this.getIf(b, this.random);
                int cIndex = b.getClassIndex();
                int mIndex = b.getMethodIndex();
                jClass = (JavaClass)jClassList.get(cIndex);
                if (jClass == null) {
                    System.out.println("error!");
                    return false;
                }

                ClassGen clsGen = new ClassGen(jClass);
                ConstantPoolGen cpg = clsGen.getConstantPool();
                Method[] methods = jClass.getMethods();
                MethodGen mthGen = new MethodGen(methods[mIndex], b.getClassName(), cpg);
                InstructionList ifInstr = mthGen.getInstructionList();
                ifInstr.setPositions();
                InstructionHandle ifHandle = ifInstr.findHandle(ifIndex);
                if (ifHandle != null && ifHandle.getInstruction() instanceof IfInstruction) {
                    IfInstruction ifIn = (IfInstruction)ifHandle.getInstruction();
                    InstructionHandle target = ifIn.getTarget();
                    Predicates pre = new Predicates();
                    new ArrayList();
                    ArrayList var = this.findLocalVariable(ifInstr, cpg, ifHandle);
                    if (var.size() != 0) {
                        InstructionList ins = pre.createOpaqueInsts((Integer)watermark.get(i++), var);
                        System.out.println("class" + clsGen.getClassName() + " method:" + mthGen.getName());
                        InstructionHandle h = ins.getEnd();
                        Instruction lastInst = h.getInstruction();
                        if (lastInst instanceof IfInstruction) {
                            IfInstruction lastIf = (IfInstruction)lastInst;
                            lastIf.setTarget(target);
                        }

                        ifInstr.append(ifHandle, ins);
                        mthGen.setInstructionList(ifInstr);
                        mthGen.stripAttributes(false);
                        mthGen.setMaxLocals();
                        mthGen.setMaxStack();
                        mthGen.removeLocalVariables();
                        mthGen.removeAttributes();
                        mthGen.removeExceptionHandlers();
                        mthGen.removeExceptionHandlers();
                        mthGen.removeExceptions();
                        mthGen.removeLineNumbers();
                        mthGen.removeNOPs();
                        clsGen.setMethodAt(mthGen.getMethod(), mIndex);
                        jClassList.set(cIndex, clsGen.getJavaClass());
                        this.fixList(b.getIndexList(), mthGen, ifIndex);
                    }
                }
            }

            this.loadJar.outputJar(outpath, jClassList);
            return true;
        }
    }

    public ArrayList findLocalVariable(InstructionList il, ConstantPoolGen cpg, InstructionHandle ifHandle) {
        InstructionHandle[] ilh = il.getInstructionHandles();
        ArrayList al = new ArrayList();

        label44:
        for(int i = 0; i < ilh.length; ++i) {
            Instruction ins = ilh[i].getInstruction();
            if (ins instanceof LocalVariableInstruction && ilh[i].getPosition() < ifHandle.getPosition()) {
                LocalVariableInstruction lvi = (LocalVariableInstruction)ins;
                Type t = lvi.getType(cpg);
                if (t.equals(Type.INT)) {
                    for(InstructionHandle ih = ilh[i].getNext(); ih != null && ih != ifHandle; ih = ih.getNext()) {
                        if (ih.getInstruction() instanceof LocalVariableInstruction) {
                            LocalVariableInstruction another = (LocalVariableInstruction)ih.getInstruction();
                            if (another.getIndex() == lvi.getIndex()) {
                                break label44;
                            }
                        }
                    }

                    al.add(lvi);
                }
            }
        }

        if (al.size() == 1) {
            al.add((LocalVariableInstruction)al.get(0));
        }

        return al;
    }

    public ArrayList<Integer> splitWM(String wm, String key) {
        byte[] b1 = wm.getBytes();
        byte[] b2 = key.getBytes();
        BigInteger wmInt = new BigInteger(b1);
        BigInteger keyInt = new BigInteger(b2);
        BigInteger wmKey = wmInt.add(keyInt);
        return this.split(wmKey);
    }

    public ArrayList<Integer> split(BigInteger wmKey) {
        int maxValue = 1000;
        BigInteger comb = BigInteger.ONE;

        int n;
        for(n = 0; wmKey.compareTo(comb) >= 0; comb = comb.multiply(BigInteger.valueOf((long)(n + maxValue))).divide(BigInteger.valueOf((long)n))) {
            wmKey = wmKey.subtract(comb);
            ++n;
        }

        ArrayList<Integer> parts = new ArrayList();
        if (n == 0) {
            return parts;
        } else {
            Fenjie ma = new Fenjie();
            SortedSet<Integer> combValue = ma.getComb(wmKey, n + maxValue, maxValue);
            Iterator i = combValue.iterator();
            int prev = -1;

            int curr;
            for(int currValue = 0; i.hasNext(); prev = curr) {
                if (currValue > maxValue) {
                    throw new RuntimeException();
                }

                curr = (Integer)i.next();

                for(int j = 0; j < curr - prev - 1; ++j) {
                    parts.add(currValue);
                }

                ++currValue;
            }

            while(parts.size() < n) {
                parts.add(Integer.valueOf(maxValue));
            }

            return parts;
        }
    }

    public void fixList(ArrayList<Integer> indexList, MethodGen m, int oldIndex) {
        for(int i = 0; i < indexList.size(); ++i) {
            Integer iindex = (Integer)indexList.get(i);
            int index = iindex;
            if (oldIndex < index) {
                InstructionList il = m.getInstructionList();
                InstructionHandle ih = null;

                do {
                    ih = il.findHandle(index);
                    ++index;
                } while(ih == null);

                boolean notFound = true;

                while(notFound) {
                    Instruction inst = ih.getInstruction();
                    if (inst instanceof IfInstruction) {
                        Integer newIndex = new Integer(ih.getPosition());
                        indexList.set(i, newIndex);
                        notFound = false;
                    } else {
                        ih = ih.getNext();
                    }

                    if (ih == null) {
                        indexList.set(i, 0);
                        notFound = false;
                    }
                }
            }
        }

    }

    private Bundle getBundle(ArrayList candidateList, Random gen) {
        int bundleIndex = Math.abs(gen.nextInt()) % candidateList.size();
        Bundle b = (Bundle)candidateList.get(bundleIndex);
        return b;
    }

    private int getIf(Bundle b, Random gen) {
        ArrayList indexList = b.getIndexList();
        int ifRandomIndex = Math.abs(gen.nextInt()) % indexList.size();
        Integer ifIntIndex = (Integer)indexList.get(ifRandomIndex);
        int ifIndex = ifIntIndex;
        return ifIndex;
    }

    public ArrayList preprocess(ArrayList<JavaClass> classes) {
        ArrayList<Bundle> candidateList = new ArrayList();
        Bundle bundleObject = null;
        String className = null;
        JavaClass cls = null;

        for(int k = 0; k < classes.size(); ++k) {
            if (classes.get(k) != null && !((JavaClass)classes.get(k)).isAbstract() && !((JavaClass)classes.get(k)).isInterface()) {
                cls = (JavaClass)classes.get(k);
                className = cls.getClassName();
                ClassGen clg = new ClassGen(cls);
                ConstantPoolGen cpg = clg.getConstantPool();
                Method[] methods = cls.getMethods();

                for(int i = 0; i < methods.length; ++i) {
                    ArrayList<Integer> indexList = new ArrayList();
                    Method m = methods[i];
                    MethodGen mg = new MethodGen(m, className, cpg);
                    InstructionList il = mg.getInstructionList();
                    il.setPositions();
                    InstructionHandle[] ihs = il.getInstructionHandles();
                    Instruction[] insts = il.getInstructions();

                    for(int j = 0; j < insts.length; ++j) {
                        Instruction inst = insts[j];
                        if (inst instanceof IfInstruction) {
                            InstructionHandle ifHandle = ihs[j];
                            int pos = ifHandle.getPosition();
                            Integer intObj = new Integer(pos);
                            indexList.add(intObj);
                        }
                    }

                    if (indexList.size() > 0) {
                        bundleObject = new Bundle(className, m, indexList, k, i);
                        candidateList.add(bundleObject);
                    }
                }
            }
        }

        if (candidateList.size() == 0) {
            this.embedError = "There are no suitable if statements to use for watermarking.";
        }

        return candidateList;
    }

    /*public String extract(String path, int key) {
        new Fenjie();
        JavaClass jClass = null;
        ArrayList<JavaClass> jClassList = this.loadJar.inputJar(path);
        this.random = new Random();
        this.random.setSeed((long)key);
        int pos = Math.abs(this.random.nextInt(jClassList.size())) % jClassList.size();
        jClass = (JavaClass)jClassList.get(pos);

        while(jClass == null) {
            jClass = (JavaClass)jClassList.get(pos);
            ++pos;
            pos %= jClassList.size();
            System.out.println(pos);
        }

        ClassGen clsGen = new ClassGen(jClass);
        ConstantPoolGen cpg = clsGen.getConstantPool();
        int i = 0;
        boolean j = false;

        ArrayList wmValue;
        for(wmValue = new ArrayList(); i < jClassList.size(); ++i) {
            jClass = (JavaClass)jClassList.get(i);
            if (jClass != null && !jClass.isAbstract() && !jClass.isInterface()) {
                clsGen = new ClassGen(jClass);
                cpg = clsGen.getConstantPool();
                Method[] methods = clsGen.getMethods();

                for (Method method : methods) {
                    if (!method.isAbstract()) {
                        MethodGen mthGen = new MethodGen(method, jClass.getClassName(), cpg);
                        InstructionList instr = mthGen.getInstructionList();
                        ArrayList<Integer> al = this.identifyLeaders(instr);
                        ArrayList blockList = this.buildBlocks(instr, al);
                        Predicates p = new Predicates();
                        if (!blockList.isEmpty()) {
                            for (int k = 0; k < blockList.size(); ++k) {
                                BasicBlock bb = (BasicBlock) blockList.get(k);
                                int temp;
                                if (bb.getLastIH().getInstruction() instanceof IfInstruction && (temp = p.isOpaque(bb.getInstList())) != -1) {
                                    wmValue.add(temp);
                                }
                            }
                        }
                    }
                }
            }
        }

        String wm;
        if (wmValue.isEmpty()) {
            wm = "no watermark";
        } else {
            wm = this.recover(wmValue, String.valueOf(key));
            wm = "the watermark is :" + wm;
        }

        return wm;
    }

    public String recover(ArrayList<Integer> wmValue, String key) {
        byte[] wmKey = key.getBytes();
        BigInteger wm = this.combine(wmValue);
        wm = wm.subtract(new BigInteger(wmKey));
        byte[] wmByte = wm.toByteArray();
        String watermark = new String(wmByte);
        return watermark;
    }

    public BigInteger combine(ArrayList<Integer> parts) {
        Fenjie m = new Fenjie();
        int maxValue = 1000;

        for(int i = 0; i < parts.size(); ++i) {
            if (maxValue < (Integer)parts.get(i) || (Integer)parts.get(i) < 0) {
                throw new IllegalArgumentException();
            }
        }

        Collections.sort(parts);
        Set<Integer> comb = new HashSet();
        int currPart = 0;
        int elem = 0;

        int n;
        for(n = 0; n < parts.size(); ++n) {
            while((Integer)parts.get(n) > currPart) {
                comb.add(new Integer(elem++));
                ++currPart;
            }

            ++elem;
        }

        while(maxValue > currPart) {
            comb.add(new Integer(elem++));
            ++currPart;
        }

        n = 0;
        BigInteger combSum = BigInteger.ZERO;

        BigInteger numer;
        BigInteger denom;
        for(BigInteger combCount = BigInteger.ONE; n < parts.size(); combCount = combCount.multiply(numer).divide(denom)) {
            combSum = combSum.add(combCount);
            ++n;
            numer = BigInteger.valueOf((long)(n + maxValue));
            denom = BigInteger.valueOf((long)n);
        }

        combSum = combSum.add(m.decodeCombination(comb, parts.size() + maxValue));
        return combSum;
    }

    public ArrayList<BasicBlock> buildBlocks(InstructionList il, ArrayList<Integer> leaders) {
        ArrayList<BasicBlock> blockList = new ArrayList();

        for(int i = 0; i < leaders.size() - 1; ++i) {
            InstructionHandle ih1 = il.findHandle((Integer)leaders.get(i));
            InstructionHandle ih2 = il.findHandle((Integer)leaders.get(i + 1));
            if (ih1.getInstruction() instanceof IfInstruction && ih2.getInstruction() instanceof IfInstruction) {
                BasicBlock bb = new BasicBlock(ih1.getNext(), ih2.getNext());
                blockList.add(bb);
            }
        }

        return blockList;
    }

    public ArrayList<Integer> identifyLeaders(InstructionList il) {
        InstructionHandle[] ihs = il.getInstructionHandles();
        ArrayList<Integer> leaders = new ArrayList();

        for(int j = 0; j < ihs.length; ++j) {
            Instruction instruction = ihs[j].getInstruction();
            if (instruction instanceof IfInstruction) {
                Integer pos = ihs[j].getPosition();
                leaders.add(pos);
            }
        }

        return leaders;
    }*/

    public static void main(String[] args) {
        OP op = new OP();
        //System.out.println(op.extract("LargeSub_fat2.jar", 22));
        //System.out.println(op.extract("jclasslib.jar", 22));
    }
}
